import os


avatar_prompts = {
    "dreamwaltz": [
        "A standing Captain Jack Sparrow from Pirates of the Caribbean",
        "Abraham Lincoln",
        "Agent 47",
        "Albert Einstein",
        "Albus Dumbledore",
        "Alice in wonderland, Disney",
        "Ariel, Disney",
        "Assassin\"s Creed",
        "Audrey Hepburn in Breakfast at Tiffany\"s",
        "Barack Obama",
        "Batman",
        "Belle, Disney",
        "Black Widow",
        "Buffy theVampire Slayer",
        "Buzz Lightyear",
        "Cao Cao",
        "Captain America",
        "Captain Marvel",
        "Carnage, Marvel character",
        "Chaplin",
        "Cinderella",
        "Darth Vader",
        "David of Michelangelo",
        "Deadpool",
        "Doctor Strange",
        "Doctor Who",
        "Donald Trump",
        "Dragon Ball",
        "Elsa",
        "Emma Stone in La La Land",
        "Eren Yeager",
        "Flynn Rider",
        "Forrest Gump",
        "Gintoki",
        "Green Arrow",
        "Harley Quinn",
        "Harry Potter, adult",
        "Hatsune Miku",
        "Hulk",
        "I am Groot",
        "Iroi Yagami",
        "Iron Man",
        "Jack Marston",
        "Jasmine, Disney",
        "Jedi",
        "Joe Biden",
        "Joker",
        "Kobe Bryant, NBA player",
        "Kratos",
        "Lara Croftin Tomb Raider",
        "Link in Zelda",
        "Lionel Messi",
        "Lord Voldemort",
        "Luke Skywalker",
        "Madonna Louise Ciccone",
        "Marie Antoinette",
        "Mario Super Mario",
        "Merida, Disney",
        "Michael Jackson",
        "Michael Jordan",
        "Mobile Suit Gundam",
        "Monkey D. Luffy",
        "Mulan",
        "Napoleon",
        "Naruto Uzumaki",
        "Neo in the Matrix",
        "Nick Fury",
        "Nick Wilde",
        "Optimus Prime",
        "Peter Pan",
        "Pinocchio",
        "Princess Leia",
        "Pudge",
        "Queen Elizabeth II",
        "Rapunzel",
        "Retro astronaut with a helmet and a bulky suit",
        "Robot",
        "Ronald Weasley",
        "Rose in Titanic",
        "Runpunzelin Tangled",
        "Saber in Fate",
        "Sailor Moon",
        "Santa Claus",
        "Sherlock Holmes",
        "Shinichi Kudo in Case Closed",
        "Snow White, Disney",
        "Soldier",
        "Spiderman",
        "Stormtrooper",
        "Sun Wukong",
        "Superman",
        "Tarzan",
        "Taylor Swift",
        "Teletubbies",
        "Thanos",
        "The Godfather",
        "Thor",
        "Tiger Woods",
        "Tinker Bell",
        "Ultraman",
        "Venus de Milo",
        "Winston Churchill",
        "Wolverine, marvel character",
        "Wonder Woman",
        "Woody in Toy Story",
        "Yoda",
        "ballet dancer",
        "baseball player",
        "blue fairy with wings",
        "chimpanzee",
        "doctor wearing stethoscope",
        "football player",
        "hockey player",
        "king of pandaren",
        "ping pong player",
    ],

    "chatgpt": [
        "a chef in a classic white apron",
        "a chef in a crisp white chef\"s coat and toque",
        "a child in a superhero costume with a red cape",
        "a college student in a hoodie and backpack",
        "a firefighter in full gear, holding an axe",
        "a firefighter in full protective gear",
        "a gardener in overalls and a wide-brimmed hat",
        "a hiker in rugged outdoor gear",
        "a lifeguard in a red swimsuit",
        "a man in workout attire, headphones in and a gym bag slung over his shoulde",
        "a musician in a leather jacket",
        "a nurse in scrubs and a stethoscope",
        "a scientist in a lab coat and safety goggles",
        "a skateboarder in baggy jeans and a graphic tee",
        "a street musician with a guitar and worn jeans",
        "a teenager with a skateboard, wearing torn jeans and a beanie",
        "a teenager with colorful tattoos and piercings",
        "a woman in a crisp business suit, carrying a leather briefcase",
        "a woman in a floral sundress",
        "a yoga instructor in comfortable athleisure wear",
        "a young child in a superhero costume",
        "a young man in a sharp business suit",
        "a young woman with a backpack and hiking boots",
        "a young woman with curly hair in a red sundress",
        "an artist covered in paint",
        "an elderly gentleman in a tweed jacket and bowtie",
        "an elderly woman with a floral dress and a sunhat",
        "an elderly woman with a flowing floral dress",
    ],

    "creative": [
        "a boxer wearing swimsuit",
        "a chef in a lab coat and safety goggles",
        "a clown in a superhero costume",
        "a doctor wearing sunhat holding flowers",
        "a lifeguard wearing business suit",
        "a little chubby boy in a sharp business suit",
    ],

	"dreamhuman": [
        "a Black female surgeon",
        "a Black man wearing a green t-shirt",
        "a Black man wearing a red baseball cap",
        "a Black person in a military uniform",
        "a Black woman dressed in gym clothes",
        "a Black woman in a wedding dress",
        "a Black woman wearing a hoodie",
        "a Buddhist monk",
        "a Medieval European King",
        "a Mediterranean man with a beard wearing a white linen shirt",
        "a Roman soldier wearing his armor",
        "a Spanish flamenco dancer",
        "a Viking",
        "a ballerina",
        "a bedouin dressed in white",
        "a blonde woman wearing yoga pants",
        "a body builder wearing a tanktop",
        "a bodybuilder wearing a tanktop",
        "a boxer",
        "a chef dressed in white",
        "a clown",
        "a farmer",
        "a female professor wearing a regalia",
        "a high-school janitor",
        "a karate master wearing a Black belt",
        "a male gymnastics athlete",
        "a man wearing a Christmas sweater",
        "a man wearing a Hawaiian shirt, sunglasses and shorts",
        "a man wearing a bomber jacket",
        "a man wearing a hoodie",
        "a man wearing a striped shirt and white linen pants",
        "a man with Black hair wearing a vest",
        "a man with blond hair wearing a brown leather jacket",
        "a man with dreadlocks",
        "a member of a rock band",
        "a ninja",
        "a person dressed at the Venice Carnival",
        "a person in a diving suit",
        "a plus-size model wearing pyjamas",
        "a policewoman",
        "a pregnant person of color",
        "a security guard",
        "a senior Black person wearing a polo shirt",
        "a slim man wearing a navy blazer and gray trousers",
        "a thin Marathon runner",
        "a track and field athlete",
        "a woman in traditional Bavarian clothing",
        "a woman wearing a long dress",
        "a woman wearing a short jean skirt and a cropped top",
        "a woman wearing a wedding dress",
        "a woman wearing ski clothes",
        "a woman with long blond hair",
        "a young man wearing a turtleneck",
        "an African woman dressed in traditional clothes",
        "an Asian man wearing a navy suit",
        "an Indian bride in a traditional dress",
        "an elderly man wearing a beige suit",
        #---------- Figure 7 ----------#
        # "astronaut",
        # "construction manager",
        # "firefighter",
        # "gardener",
        # "pilot",
        # "police officer",
        # "robot",
        # "senior citizen",
        # "soldier",
        # "teenager",
        # "warrior",
        # "witch",
        # "wizard",
    ],

    "seeavatar": [
        "Abraham Lincoln",
        "Barack Obama",
        "Batman",
        "Captain Jack Sparrow",
        "Doctor Strange",
        "Elsa in Frozen",
        "Hillary Clinton",
        "Joker",
        "Ludwig van Beethoven",
        "Spiderman",
        "Vincent van Gogh",
        "Wolfgang Amadeus Mozart",
    ],

    "gavatar": [
        "A professional boxer",
        "Morty Smith",
        "A person in a diving suit",
        "An American soldier from World War 2",
        "Goku",
        "Rick Sanchez",
        "A person dressed at the Venice carnival",
        "A medieval European king",
        "An elderly man wearing a beige suit",
        "Kobe Bryant",
        "A man wearing a white tank top and shorts",
        "A policewoman",
        "A black female surgeon",
        "A viking",
        "Oprah Winfrey",
        "A bedouin dressed in white",
        "A farmer",
        "A clown",
        "Jane Goodall",
        "Homer Simpson",
        "Kristoff in Frozen",
        "Luffy in one piece",
        "Spiderman",
        "Jeff Bezos",
    ],
}


def read_txt_file(txt_path):
    with open(txt_path, 'r') as fio:
        lines = fio.readlines()
        characters = []
        for item in lines:
            if item.startswith('#'):
                continue
            characters.append(item.strip('\r\n ,.'))
        return characters


def get_avatar_list(set_name):
    if set_name in avatar_prompts:
        avatar_list = avatar_prompts[set_name]
        avatar_list.sort()
        return avatar_list
    else:
        assert 0, set_name


if __name__ == '__main__':
    print(avatar_prompts['gavatar'])
